from sqlalchemy import Column, String, DateTime, Text, Integer, Boolean
from sqlalchemy.ext.declarative import declarative_base
from apscheduler.jobstores.sqlalchemy import SQLAlchemyJobStore
from datetime import datetime

# 创建 SQLAlchemy 的基类
Base = declarative_base()


# 定义新的任务表结构
class CustomJob(Base):
    __tablename__ = 'custom_jobs'
    id = Column(String(191), primary_key=True)
    name = Column(String(191))  # 任务名称
    func = Column(String(191))
    trigger = Column(Text)
    args = Column(Text)
    kwargs = Column(Text)
    executor = Column(String(40))
    misfire_grace_time = Column(Integer)
    coalesce = Column(Boolean)
    max_instances = Column(Integer)
    next_run_time = Column(DateTime)
    job_state = Column(Text)
    create_time = Column(DateTime, default=datetime.now())  # 创建时间
    update_time = Column(DateTime, default=datetime.now(), onupdate=datetime.now())  # 更新时间
    params = Column(Text)  # 参数


class CustomSQLAlchemyJobStore(SQLAlchemyJobStore):
    def __init__(self, engine=None, url=None, tablename='custom_jobs', metadata=Base.metadata,
                 pickle_protocol=4, **connect_args):
        super().__init__(engine=engine, url=url, tablename=tablename, metadata=metadata,
                         pickle_protocol=pickle_protocol, **connect_args)

    def add_job(self, job):
        # 获取任务名称和参数
        name = job.kwargs.pop('name', None)
        params = job.kwargs.pop('params', None)

        # 创建新的任务记录
        custom_job = CustomJob(
            id=job.id,
            name=name,
            func='%s:%s' % (job.func.__module__, job.func.__name__),
            trigger=self.serialize_trigger(job.trigger),
            args=self.serialize_args(job.args),
            kwargs=self.serialize_args(job.kwargs),
            executor=job.executor,
            misfire_grace_time=job.misfire_grace_time,
            coalesce=job.coalesce,
            max_instances=job.max_instances,
            next_run_time=job.next_run_time,
            job_state=self.serialize_job(job),
            params=params
        )

        # 保存任务记录
        self.session.add(custom_job)
        self.session.commit()

    def update_job(self, job):
        # 获取任务名称和参数
        name = job.kwargs.pop('name', None)
        params = job.kwargs.pop('params', None)

        # 查询任务记录
        custom_job = self.session.query(CustomJob).filter_by(id=job.id).first()

        if custom_job:
            # 更新任务记录
            custom_job.name = name
            custom_job.func = '%s:%s' % (job.func.__module__, job.func.__name__)
            custom_job.trigger = self.serialize_trigger(job.trigger)
            custom_job.args = self.serialize_args(job.args)
            custom_job.kwargs = self.serialize_args(job.kwargs)
            custom_job.executor = job.executor
            custom_job.misfire_grace_time = job.misfire_grace_time
            custom_job.coalesce = job.coalesce
            custom_job.max_instances = job.max_instances
            custom_job.next_run_time = job.next_run_time
            custom_job.job_state = self.serialize_job(job)
            custom_job.params = params

            # 提交更新
            self.session.commit()

    def _reconstitute_job(self, job_state):
        job = super()._reconstitute_job(job_state)

        # 查询任务记录
        custom_job = self.session.query(CustomJob).filter_by(id=job.id).first()

        if custom_job:
            # 恢复任务名称和参数
            job.kwargs['name'] = custom_job.name
            job.kwargs['params'] = custom_job.params

        return job
